# ------------------------------------------------------------
# pyos7.py  -  The Python Operating System
#
# Added support for task sleep
# ------------------------------------------------------------

# ------------------------------------------------------------
#                       === Tasks ===
# ------------------------------------------------------------

class Task(object):
    taskid = 0
    def __init__(self,target):
        Task.taskid += 1
        self.tid     = Task.taskid   # Task ID
        self.target  = target        # Target coroutine
        self.sendval = None          # Value to send

    # Run a task until it hits the next yield statement
    @micropython.viper
    def run(self):
        return self.target.send(self.sendval)

# ------------------------------------------------------------
#                      === Scheduler ===
# ------------------------------------------------------------
from Queue import Queue
import utime as time  #  micropython


class Scheduler(object):
    def __init__(self):
        self.ready   = Queue()   
        self.taskmap = {}  

        # Task sleep
        self.sleepStartTime = {}  
        self.sleepTime = {}

        # Tasks waiting for other tasks to exit
        self.exit_waiting = {}

    def new(self,target):
        newtask = Task(target)
        self.taskmap[newtask.tid] = newtask
        self.schedule(newtask)
        return newtask.tid

    def exit(self,task):
        print ("Task {} terminated".format(task.tid))
        del self.taskmap[task.tid]
        # Notify other tasks waiting for exit
        for task in self.exit_waiting.pop(task.tid,[]):
            self.schedule(task)

    def waitforexit(self,task,waittid):
        if waittid in self.taskmap:
            self.exit_waiting.setdefault(waittid,[]).append(task)
            return True
        else:
            return False
    def sleep(self, task, sleepTime):
        if task.tid in self.taskmap:
            self.sleepStartTime[task.tid] = time.ticks_ms()
            self.sleepTime[task.tid] = sleepTime
            self.schedule(task)
            return True
        else:
            return False
    @micropython.viper
    def schedule(self,task):
        self.ready.put(task)
    @micropython.viper
    def mainloop(self):
         while self.taskmap:
            task = self.ready.get()
            if task.tid in self.sleepStartTime:
                if(abs((time.ticks_ms()-self.sleepStartTime[task.tid]))>=self.sleepTime[task.tid]):
                    del self.sleepStartTime[task.tid]
                    del self.sleepTime[task.tid]
                else:
                    self.schedule(task)
                    continue
            try:
                result = task.run()
                if isinstance(result,SystemCall):
                    result.task  = task
                    result.sched = self
                    result.handle()
                    continue
            except StopIteration:
                self.exit(task)
                continue
            self.schedule(task)

# ------------------------------------------------------------
#                   === System Calls ===
# ------------------------------------------------------------

class SystemCall(object):
    def handle(self):
        pass

# Return a task's ID number
class GetTid(SystemCall):
    def handle(self):
        self.task.sendval = self.task.tid
        self.sched.schedule(self.task)

# Create a new task
class NewTask(SystemCall):
    def __init__(self,target):
        self.target = target
    def handle(self):
        tid = self.sched.new(self.target)
        self.task.sendval = tid
        self.sched.schedule(self.task)

# Kill a task
class KillTask(SystemCall):
    def __init__(self,tid):
        self.tid = tid
    def handle(self):
        task = self.sched.taskmap.get(self.tid,None)
        if task:
            task.target.close() 
            self.task.sendval = True
        else:
            self.task.sendval = False
        self.sched.schedule(self.task)

# Wait for a task to exit
class WaitTask(SystemCall):
    def __init__(self,tid):
        self.tid = tid
    def handle(self):
        result = self.sched.waitforexit(self.task,self.tid)
        self.task.sendval = result
        # If waiting for a non-existent task,
        # return immediately without waiting
        if not result:
            self.sched.schedule(self.task)

# Task to sleep(ms) 
class SleepTask(SystemCall):
    def __init__(self, sleepTime):
        self.sleepTime = sleepTime
    def handle(self):
        result = self.sched.sleep(self.task,self.sleepTime)
        
# ------------------------------------------------------------
#                      === Example ===
# ------------------------------------------------------------
if __name__ == '__main__':
    def foo():
        tid = yield GetTid()
        #for i in range(0, 5):
        while True:
            print ("I'm foo", tid)
            yield SleepTask(1000)
    def foo2():
        tid = yield GetTid()
        #for i in range(0, 5):
        while True:
            print ("I'm foo", tid)
            yield SleepTask(2000)

    def main():
        child = yield NewTask(foo())
        print ("Waiting for child")
        yield NewTask(foo2())
        yield WaitTask(child)
        print ("Child done")
        print(time.time())

    sched = Scheduler()
    sched.new(main())
    sched.mainloop()
